package pogvue.gui;

import pogvue.analysis.AAFrequency;
import pogvue.datamodel.Alignment;
import pogvue.datamodel.AlignmentI;
import pogvue.datamodel.SequenceI;
import pogvue.datamodel.SequenceNode;
import pogvue.io.FileParse;
import pogvue.io.FastaFile;

import javax.swing.*;
import java.awt.*;
import java.awt.event.MouseEvent;
import java.awt.event.MouseListener;
import java.awt.event.MouseMotionListener;
import java.awt.geom.AffineTransform;
import java.io.IOException;
import java.util.Enumeration;
import java.util.Hashtable;
import java.util.Vector;

public final class KmerScore  extends JPanel implements MouseListener,MouseMotionListener  {
  private CountPanel cp;
  private CountPanel cp2;


  Vector counts;
  private final int    num;

  private double[][] score;
  private double[]   max;

  private final Vector kmers;
  private final Vector conserved;

  private double totmax = 0.0;
  private double totmin = 0.0;

  final Font font = new Font("Courier",Font.PLAIN,8);

  private int scale = 8;

  private int xprev;
  private int yprev;

  private KmerScore(Vector conserved, int num,Vector kmers,int scale) {
    this.conserved = conserved;
    this.num    = num;
    this.kmers  = kmers;
    this.scale  = scale;

    addMouseListener(this);
    addMouseMotionListener(this);
  }

  private void score() {

     score = new double[num][num];     // Conditional probabilities for each kmer
     max   = new double[num];

     int totpairs = 0;
     int totkmers = 0;

     double probs[] = new double[num]; // Probabilities for each kmer

     for (int i = 0; i < num; i++) {
       probs[i] = 0.0;
       max[i] = 0.0;
       for (int j = 0; j < num; j++) {
           score[i][j] = 0.0;
       }
     }
        
     System.out.println("Num " + num);

     for (int i = 0;i < conserved.size(); i++) {

          //This is one rows worth of conserved positions
          Hashtable tmp = (Hashtable)conserved.elementAt(i);

          Enumeration en = tmp.keys();

          // We loop over each kmer
          while (en.hasMoreElements()) {
              String key = (String)en.nextElement();
              Vector pos = (Vector)tmp.get(key);

              totkmers += pos.size();
              probs[kmers.indexOf(key)]+= pos.size(); 

              // We then loop over each position for that kmer

	        for (int j = 0; j < pos.size(); j++) {
                    int pp = (Integer) pos.elementAt(j);
          
                    Enumeration en2 = tmp.keys();

                    while (en2.hasMoreElements()) {
                        String k = (String)en2.nextElement();

                        if (!k.equals(key)) {
                            Vector pos2 = (Vector)tmp.get(k);

                            for (int kk = 0; kk < pos2.size(); kk++) {

                                  int p2 = (Integer) pos2.elementAt(kk);

                                  if (Math.abs(pp - p2) > 6) {
                                     score[kmers.indexOf(key)][kmers.indexOf(k)]++; 
                                     score[kmers.indexOf(k)][kmers.indexOf(key)]++; 
                                     totpairs+=2;
                                 }
                            }
                       }
                   }
              }
          }
     }


     // score [i][j] contains the number of instances when i and j are present     
     // It should be symmetric


     // Convert the kmer i counts into probabilities

     for (int i = 0; i < num; i++) {
         probs[i] /= (1.0*totkmers);

         System.out.println("Probs " + i + " " + kmers.elementAt(i) + " " + probs[i]);
     }

     totmax = 0.0;

     // Now convert the symmetric score[i][j] into condtional probabilities on j
     // and also divide by probability of i

     for (int i = 0; i < num; i++) {

         for (int j = 0; j < num; j++) {

            score[i][j] /= (1.0*totpairs);
            score[i][j] /= probs[j];
            score[i][j] /= probs[i];

            System.out.print(score[i][j] + "\t");

            if (j != i) {
              if (score[i][j] > totmax) { totmax = score[i][j];}

            }
          }
          System.out.println();
     }

      Vector s = new Vector();

      // Sort columns first - cp has column positions
      for (int ii = 0; ii < num; ii++) {
           Vector s2 = new Vector();

           for (int k = 0;  k < num; k++) {
               s2.addElement((int) (score[k][ii] * 1000));
           }
           s.addElement(s2);
      }

      cp = new CountPanel(s,num,num);

      cp.cluster();

      Vector s3 = new Vector();

      // Now sort rows - cp2 has rows

      for (int ii = 0; ii < num; ii++) {
           Vector s4 = new Vector();

           for (int k = 0;  k < num; k++) {
               s4.addElement((int) (score[ii][k] * 1000));
           }
           s3.addElement(s4);
      }

      cp2 = new CountPanel(s3,num,num);

      cp2.cluster();

      totmax = 0.0;
      totmin = 0.0;
  
      for (int i = 0; i < num; i++) {

          String str = ((SequenceNode)cp.v.elementAt(i)).getName();
          int nn = Integer.parseInt(str);
          String name = (String)kmers.elementAt(nn);
          System.out.print(name + "\t");
      }

      System.out.println();

      for (int i = 0; i < num; i++) {
          String na   = ((SequenceNode)cp2.v.elementAt(i)).getName();
          int    nn   = Integer.parseInt(na);
          String name = (String)kmers.elementAt(nn);

          System.out.print("NewScore\t" + name + "\t");

	  for (int j = 0; j < num; j++) {
	      String str2 = ((SequenceNode)cp.v.elementAt(j)).getName();
	      int n = Integer.parseInt(str2);

	      int ii = n;
	      int jj = nn;

	      System.out.print((int)(1000*score[ii][jj]) + "\t");

	      if (score[ii][jj] < 1e-10) {
		  score[ii][jj] = Math.log(1e-10)/Math.log(2);
	      } else {
		  score[ii][jj] = Math.log(score[ii][jj])/Math.log(2);
	      }
	      
	      
	      if (score[ii][jj] > totmax) {
		  totmax = score[ii][jj];
	      }
	      if (score[i][j] < totmin) {
		  totmin = score[ii][jj];
	      }
	      if (score[ii][jj] > max[ii]) {
		  max[ii] = score[ii][jj];
	      }
	  }
	  System.out.println();
	  
      }

      totmin *= -1.0;

      System.out.println("Max/min " + totmax + " " + totmin);

  }

public void paint(Graphics g){

     g.setColor(Color.white);
     g.fillRect(0,0,size().width,size().height); 

     for (int i = 0; i < num; i++) {

       // cp has columns
       String name2 = ((SequenceNode)cp.v.elementAt(i)).getName();

       int n2 = Integer.parseInt(name2);

       for (int j = 0; j < num; j++) {

          // cp2 has rows 
          String name = ((SequenceNode)cp2.v.elementAt(j)).getName();
          int n = Integer.parseInt(name);

          if (n != n2) {

            if (score[n][n2] > 0) {

              g.setColor(new Color((int)(255.0*score[n][n2]/totmax),0,0));

            } else {

              int colpos = (int)(-255.0*score[n][n2]/(totmin));

              g.setColor(new Color(0,colpos,colpos));
            }
            g.fillRect(j*scale,i*scale,scale,scale);
          }
       }
     }

    Graphics2D g2d = (Graphics2D)g;
    AffineTransform at = new AffineTransform();


    g2d.setColor(Color.black);

    for (int i = 0; i < num; i++) {

      // cp2 rows
      String name2 = ((SequenceNode)cp2.v.elementAt(i)).getName();

      int n2 = Integer.parseInt(name2);

      g2d.drawString((String)kmers.elementAt(n2),num*scale+4,i*scale+scale);

    }

    for (int j = 0; j < num; j++) {
      String name = ((SequenceNode)cp.v.elementAt(j)).getName();
      int n = Integer.parseInt(name);
      int k = 0;
      int l = ((String)kmers.elementAt(n)).length();
      while (k < l) {
         g2d.drawString(((String)kmers.elementAt(n)).substring(k,k+1),j*scale,num*scale+(1+k)*6);
         k++;
      }
    }
  }

  public static void main(String[] args) {
    try {

        FileParse fp = new FileParse(args[0],"File");
			  FastaFile kf = new FastaFile(args[1],"File");
			  SequenceI[] kmer = kf.getSeqsAsArray();

        System.out.println("Kmer " + kmer.length);

        Vector conserved = new Vector();

        String line;

        while ((line = fp.nextLine()) != null) {
					FastaFile tf = new FastaFile("Users/mclamp/cvs/mut_scripts/tss3/tot/"+line,"File");

					SequenceI[] s = tf.getSeqsAsArray();

          Hashtable conshash = conservedKmers(new Alignment(s),kmer);
          conserved.addElement(conshash);
          Enumeration en = conshash.keys();

          while (en.hasMoreElements()) {
            String key = (String)en.nextElement();
 
            Vector v = (Vector)conshash.get(key);
            if (v.size() > 0) {
              System.out.println("Vec " + key + " " + v);
            }

        }
      }

      Vector kmers = new Vector();

			for (SequenceI aKmer : kmer) {
				kmers.add(aKmer.getSequence());
			}
			int scale = Integer.parseInt(args[2]);

      KmerScore ks = new KmerScore(conserved,kmer.length,kmers,scale);

      ks.score();
            
      JFrame jf = new JFrame();
      jf.getContentPane().add(ks);

   
      jf.setSize(scale*kmer.length+40,scale*kmer.length+40);
      jf.show(); 
   } catch (IOException e) {
     System.out.println("Can't read file");
  }
  }

  public void mouseEntered(MouseEvent evt) { }
  public void mouseExited(MouseEvent evt)  { }

  public void mouseClicked(MouseEvent evt) {
      int x = evt.getX()/scale;
      int y = evt.getY()/scale;

      // cp2 rows
      int n1 = Integer.parseInt(((SequenceNode)cp2.v.elementAt(x)).getName());

      // cp columns
      int n2 = Integer.parseInt(((SequenceNode)cp.v.elementAt(y)).getName());

      String kmer1 = (String)kmers.elementAt(n1);
      String kmer2 = (String)kmers.elementAt(n2);

      System.out.println("X Y " + x + " " + y + " " + kmer1 + " " + kmer2 + " " + score[n1][n2]);

  }

  public void mouseReleased(MouseEvent evt) {}
  public void mousePressed(MouseEvent evt) {}
  public void mouseMoved(MouseEvent evt) { }
  public void mouseDragged(MouseEvent evt) {
      int x = evt.getX()/scale;
      int y = evt.getY()/scale;

      if (x != xprev || y != yprev) {
      int n1 = Integer.parseInt(((SequenceNode)cp2.v.elementAt(x)).getName());
      int n2 = Integer.parseInt(((SequenceNode)cp.v.elementAt(y)).getName());

      String kmer1 = (String)kmers.elementAt(n1);
      String kmer2 = (String)kmers.elementAt(n2);

      System.out.println("X Y " + x + " " + y + " " + kmer1 + " " + kmer2);
      }
      xprev = x;
      yprev = y;
  }

  private static Hashtable conservedKmers(AlignmentI align,SequenceI[] kmers) {

      Hashtable conshash = new Hashtable();

      // initialize conserved to 0 for all kmers

		for (SequenceI kmer1 : kmers) {
			conshash.put(kmer1.getSequence(), new Vector());
		}

		Vector sel = new Vector();

      for (int i = 0; i < align.getHeight(); i++) {
         if (align.getSequenceAt(i).getName().indexOf("Ensembl") >= 0) {
            sel.addElement(align.getSequenceAt(i));
         }
      }
      // find all blocks in the sequence

      Hashtable blocks = AAFrequency.findBlockStarts(align.getSequences(),
                                                  1,
                                                  align.getWidth(),
                                                  sel);


      SequenceI ref = align.getSequenceAt(0);

      // loop over all kmers
		for (SequenceI kmer2 : kmers) {
			String kmer = kmer2.getSequence();

			//System.out.println("Kmer " + kmer);
			// foreach block
			Enumeration en = blocks.keys();

			while (en.hasMoreElements()) {
				Integer pos = (Integer) en.nextElement();
				int len = (Integer) blocks.get(pos);

				//System.out.println("Pos/len " + pos + " " + len);

				if (len >= kmer.length() && ref.getLength() > (pos + len)) {
					int start = pos;

					String bstr = ref.getSequence().substring(pos, pos + len + 1);

					// is this kmer conserved - add one to count if so
					if (bstr.indexOf(kmer) >= 0) {

						System.out.println("Bstr " + start + " " + kmer + " " + bstr + " " + (pos + bstr.indexOf(kmer)));
						Vector v = (Vector) conshash.get(kmer);
						v.addElement(pos.intValue() + bstr.indexOf(kmer));
						//  System.out.println("Adding " + kmer + " " + ((pos.intValue()+bstr.indexOf(kmer))));
					}

				}
			}
		}
		return conshash;
  }


}


